import numpy as np
import scipy.stats as stat
from scipy.optimize import fsolve


def getNormVal(mean, sigma2, size, sampleNum, newDistMaxX):
    result = np.zeros((size,))
    toMovedX = lambda x: x - (newDistMaxX - mean)

    i = 0
    print('t: ', newDistMaxX)
    while i < size:
        # print('before:', i, end=' ')
        iVal = toMovedX(i)
        # print('after:', iVal)
        cdfSub = stat.norm.cdf(iVal + 0.5, loc=mean, scale=sigma2) - stat.norm.cdf(iVal - 0.5, loc=mean, scale=sigma2)
        result[i] = cdfSub * sampleNum
        i += 1

    result *= sampleNum / result.sum()
    return result


def toNorm(dist, newDistMaxX=None):
    size = len(dist)
    maxX = np.argmax(dist)
    sampleNum = np.sum(dist)

    sigma2 = np.std(dist) ** 2
    sigma2 = -0.136 * sigma2 + 6.578

    if newDistMaxX is None:
        newDistMaxX = maxX

    return getNormVal(0, sigma2, size, sampleNum, newDistMaxX)


def toUniform(dist2, newDistMaxX=None):
    n = len(dist2)
    avg = np.average(dist2)
    var = np.var(dist2)

    def equ(i):
        x = i[0]
        y = i[1]
        return [(x + (n - 1) * y) / n - avg,
                ((((n - 1) * (y - avg) ** 2) + (x - avg) ** 2) / n) - var]

    r = fsolve(equ, np.array([0, 0]))

    if newDistMaxX is None:
        newDistMaxX = np.argmax(dist2)
    dist2_true_norm = np.full((n,), r[1])
    dist2_true_norm[newDistMaxX] = r[0]
    dist2_true_norm *= sum(dist2) / dist2_true_norm.sum()
    return dist2_true_norm


if __name__ == '__main__':
    x =  [1, 2, 2, 1]
    print(toNorm(x, 3))
    print(toNorm(x, 1))
